from .base import pool, cates, BaseGrouping
from models import Sample
from typing import List, Set
from concurrent.futures import Future
import os


class DirectoryGroupingEngine(BaseGrouping):
    task_record: List[Future] = []

    def __init__(self, export_location: str, category_file: str = "category", sample_file: str = "sample",
                 validation_file: str = "validation", resize_to: int = 0, padding: bool = False,
                 sample_type: str = "png", categories: Set[str] = None, group: int = 1, training: float = 0.,
                 validate: float = 0.):
        super().__init__(export_location, category_file, sample_file, validation_file,
                         resize_to, padding, sample_type, categories, group, training, validate)

    def calculate_works(self) -> int:
        total_samples = Sample.select() \
            .from_(Sample, cates) \
            .where(cates.c.value << self.categories) \
            .count()
        return total_samples

    def split_group(self):
        import random
        self.groups.clear()
        categoried_sample = dict()
        for category in self.categories:
            samples = Sample.select(Sample).from_(Sample, cates).where(cates.c.value == category)
            categoried_sample.update({category: {s.target_fingerprint for s in samples.objects()}})
        for category in self.categories:
            group_size = len(categoried_sample[category]) // self.split_groups
            for i in range(self.split_groups):
                chooses = set(random.sample(categoried_sample[category], group_size))
                categoried_sample.update({category: categoried_sample[category] - chooses})
                training_size = round(len(chooses) * self.training_ratio)
                training_group = set(random.sample(chooses, training_size))
                try:
                    self.groups[i].update({category: (training_group, chooses - training_group)})
                except IndexError:
                    self.groups.append({category: (training_group, chooses - training_group)})

    def generate_samples(self):
        for i, group in enumerate(self.groups):
            training_path = os.path.join(self.export_location, f"group{i:03d}", self.sample_file)
            validate_path = os.path.join(self.export_location, f"group{i:03d}", self.validation_file)
            if not os.path.exists(training_path):
                os.makedirs(training_path)
            if not os.path.exists(validate_path):
                os.makedirs(validate_path)
            for category in group:
                export_category = (os.path.join(training_path, category),
                                   os.path.join(validate_path, category))
                for p in export_category:
                    if not os.path.exists(p):
                        os.makedirs(p)
                for g in range(2):
                    for fp in group[category][g]:
                        sample = Sample.get(Sample.target_fingerprint == fp)
                        self.task_record.append(
                            pool.submit(
                                self.save_sample_picture,
                                export_category[g],
                                sample
                            )
                        )

    def cancel(self):
        for task in self.task_record:
            task.cancel()
